
#include <assert.h>
#include <autoconf.h>
#include <libut.h>
#include <nfs/nfs.h>
#include <string.h>

#include "../macros.h"
#include "../mapping.h"
#include "../swaptable.h"
#include "../task.h"
#include "../ut_manager/ut.h"
#include "../vmem_layout.h"

#include "../frametable.h"

#include "../debug.h"
#define verbose SOS_DEBUG_FRAMETABLE
#include <sys/debug.h>
#include <sys/panic.h>

/* Leak frames from UT for internal use - i.e. allocate frame table */
int _range_alloc_retype_and_map(seL4_Word vaddr, seL4_Word size);
static seL4_Word _frame_retype_and_map(seL4_Word paddr, seL4_Word vaddr);

/* Indexing functions.
   A frame_table_entry is a cell in the _frame_table array.
   A frame_number is the fte's index in that array. Pager uses these.
   We map all allocated frames into SOS, and call these addresses sos_vaddr.

   The frame_number is used to obtain both the fte* and the sos_vaddr.
            frame_table_entry <-> frame_number <-> sos_vaddr
 */
frame_table_entry *_frame_number_to_fte(seL4_Word frame_num);
seL4_Word _fte_to_framenumber(frame_table_entry *fte);
frame_table_entry *_vaddr_to_fte(seL4_Word sos_vaddr);

static frame_table _frame_table;

seL4_Word _sys_frames;     /* Total frames in system     */
seL4_Word _max_frames;     /* Max frames managable       */
seL4_Word _cur_frames = 0; /* # frames alloced right now */

static seL4_Word _ut_hi;
static seL4_Word _ut_lo;

extern frame_table_entry *_swap_queue;

int frame_table_init(void) {
    int err;
    uint32_t table_size;

    // Die on double init
    conditional_panic((_frame_table != NULL), "Frame table already initialised");

    // Sanity-check our layout
    assert(FRAME_TABLE_START > DMA_VEND);
    assert(PAGE_IS_ALIGNED(FRAME_TABLE_START));

    // Figure out how much addressable physical memory we have in ut.
    // We manage all the way from 0 to _ut_hi, for easy addressing.
    ut_find_memory(&_ut_lo, &_ut_hi);

    // Offset max frames by ut_lo. Recall that the table covers [0-ut_lo] too.
    _max_frames = PHYS_MAX_FRAMES;
    _sys_frames = DIVROUND(_ut_hi, PAGE_SIZE);
    table_size = _sys_frames * sizeof(struct frame_table_entry);

    // Frame table is bounded, to avoid running into other mapped areas.
    // XXX Should be able to proceed with reduced-size frame table...
    if (table_size > FRAME_TABLE_SIZE) {
        dprintf(1, "FATAL: too much physical mem. Rewrite the frame table.\n");
        assert(0);
    }

    dprintf(2, "%d frames of addressable physical memory (%d bytes)\n",
            _sys_frames, PAGE_UNSHIFT(_sys_frames));
    dprintf(2, "Frame table occupies %d frames (%dB)\n",
            DIVROUND(table_size, PAGE_SIZE), table_size);

    // Init swap file here to avoid races / double-initialisation
    _swap_file_init(NULL, NULL);

    // Grab a range and map it contiguously outside of the window
    err = _range_alloc_retype_and_map(FRAME_TABLE_START, table_size);
    if (err) {
        dprintf(1, "Failed to reserve range for frame table\n");
        return FRAME_TABLE_OOM;
    } else {
        // We are relying on UT giving us clean frames.
        _frame_table = (frame_table) FRAME_TABLE_START;
        return FRAME_TABLE_SUCCESS;
    }
}

// Since frame_alloc can be either synchronous OR asynchronous, we need to let the
// calling function return. i.e. we present a 100% asynchronous interface.
// We do this by registering a high-priority SOS task to perform the callback.
typedef struct frame_alloc_cb_args {
    seL4_Word frame_num;
    seL4_Word vaddr;
    int err;
    frame_alloc_cb cb;
    void *data;
} frame_alloc_cb_args;

static int _alloc_cb_task(void *cookie) {
    frame_alloc_cb_args *args = (frame_alloc_cb_args *) cookie;
    if (args->err) {
        args->vaddr = 0;
    } else {
        args->vaddr = frame_translate_number(args->frame_num);
        dprintf(8, "frame_alloc: clearing frame, calling back (%d @ %p)\n",
                args->frame_num, PAGE_ALIGN(args->vaddr));
        // Clear the frame
        bzero((void *) PAGE_ALIGN((args->vaddr)), PAGE_SIZE);
        // Flush the instruction cache
        frame_flush_icache(args->frame_num);
    }
    // Call back
    dprintf(8, "frame_alloc_cb: args->vaddr = %p\n", args->vaddr);
    (*(args->cb))(args->vaddr, args->data);
    free(args);
    dprintf(8, "frame_alloc: done calling back\n");
    return SOS_TASK_SUCCESS;
}

static void _frame_swapped_out(int err, seL4_Word frame_number,
                               seL4_Word swap_number, void *cookie) {
    if (err == FRAME_TABLE_SWAPRACE) {
        dprintf(6, "frame_swapped_out: lost race, trying again\n");
        frame_swapout(&_frame_swapped_out, cookie);
    } else {
        frame_alloc_cb_args *data = (frame_alloc_cb_args *) cookie;
        data->frame_num = frame_number;
        data->err = err;
        sos_task_add_ready(&_alloc_cb_task, cookie,
                           SOS_TASK_PRIORITY_HIGH);
    }
    return;
}

int frame_alloc(frame_alloc_cb cb, void *data) {
    int err;
    seL4_Word paddr;
    seL4_CPtr sos_cap;
    frame_table_entry *fte;

    // Die if uninitialised
    conditional_panic((_frame_table == NULL), "Frame table uninitialised");
    // Can't proceed without a callback
    conditional_panic(!cb, "Invalid pointer cb");

    // All the information needed to perform the callback
    frame_alloc_cb_args *args = malloc(sizeof(struct frame_alloc_cb_args));
    if (args == NULL) {
        dprintf(1, "frame_alloc: Could not allocate space for callback\n");
        // XXX Use scratch pool
        return FRAME_TABLE_OOM;
    }
    bzero((void *) args, sizeof(struct frame_alloc_cb_args));

    // Pack useful information into the heap
    args->cb = cb;
    args->data = data;

    if (_cur_frames >= _max_frames) {
        // Out of memory - hit artificial limit
        dprintf(6, "frame_alloc: hit artificial frame limit. swapping...\n");
        err = frame_swapout(&_frame_swapped_out, (void *) args);
        if (err) {
            // Error out and call back - failed to swap out
            args->err = err;
            sos_task_add_ready(&_alloc_cb_task, (void *) args,
                               SOS_TASK_PRIORITY_HIGH);
            return args->err;
        }
    } else {
        // Get frame from ut_alloc
        paddr = ut_alloc(PAGE_BITS);
        if (paddr == 0) {
            // Out of memory - hit physical limit.
            err = frame_swapout(&_frame_swapped_out, (void *) args);
            if (err) {
                // Error out and call back - failed to swap out
                args->err = err;
                sos_task_add_ready(&_alloc_cb_task, (void *) args,
                                   SOS_TASK_PRIORITY_HIGH);
                return args->err;
            }
        } else {
            // Got a valid frame from UT.
            _cur_frames++;

            // Figure out our vaddr within the SOS window - just add FRAME_TABLE_START.
            args->vaddr = WINDOW_START + paddr;
            // assert(args->vaddr < (WINDOW_END - (PAGE_SIZE - 1)));
            args->frame_num = PAGE_SHIFT(paddr);

            // Retype as SmallPageObject and map into SOS window
            sos_cap = _frame_retype_and_map(paddr, args->vaddr);

            // Set elements in _frame_table using shifted paddr as index
            fte = _frame_table + args->frame_num;
            fte->page = (seL4_ARM_Page) sos_cap;

            // Set up callback
            args->err = FRAME_TABLE_SUCCESS;
            err = sos_task_add_ready(&_alloc_cb_task, (void *) args,
                                     SOS_TASK_PRIORITY_HIGH);
            if (err) {
                dprintf(1, "sos_frame_alloc: Could not register callback (%d)\n", err);
            }
        }
    }
    dprintf(6, "frame_alloc: sync call finished, cb pending\n");
    return FRAME_TABLE_SUCCESS;
}

int frame_grant(seL4_Word sos_vaddr, sos_pcb *pcb, seL4_Word proc_vaddr) {
    uint32_t frame_number;
    frame_table_entry *fte;

    // Check that frame table is initialised
    conditional_panic((_frame_table == NULL), "Frame table uninitialised");

    // Check that the cspace is valid
    conditional_panic((pcb == NULL), "Process cspace does not exist");

    // Check that the frame is allocated / tracked
    frame_number = frame_translate_vaddr(sos_vaddr);
    fte = _frame_table + frame_number;
    if (!(fte->page) || fte->mapping.proc_cap) {
        return FRAME_TABLE_INVAL;
    }

    // Break out if already granted - an attempted frame-share!
    if (fte->mapping.proc_cap != 0) {
        dprintf(6, "frame_grant: frame has already been granted\n");
        return FRAME_TABLE_SHARE_FAULT;
    }

    // cspace cap copy
    fte->mapping.proc_cap = cspace_copy_cap(cur_cspace, cur_cspace, fte->page,
                                            seL4_AllRights);
    fte->mapping.proc_vaddr = proc_vaddr;
    fte->mapping.pcb = pcb;

    // cspace copy can fail when system is OOM
    if (fte->mapping.proc_cap == CSPACE_NULL) {
        dprintf(6, "frame_grant: fauled to copy cap\n");
        return FRAME_TABLE_OOM;
    }

    // Append to the end of the swap queue
    DL_APPEND(_swap_queue, fte);
    // Swappable frame is referenced by default
    fte->swappable = 1;
    fte->referenced = 1;

    return FRAME_TABLE_SUCCESS;
}

// Pinning a frame = Making sure it can't be swapped out.
// It is sufficient to remove it from the swap queue.
int frame_pin(seL4_Word vaddr) {
    seL4_Word frame_number = frame_translate_vaddr(vaddr);
    frame_table_entry *fte = _frame_table + frame_number;

    dprintf(5, "frame_pin: pinning %p (FN %d)\n", vaddr, frame_number);

    // Delete from the swap queue, if we're in it
    if (fte->swappable) {
        DL_DELETE(_swap_queue, fte);
        fte->swappable = 0;
        fte->next = NULL;
        fte->prev = NULL;
    }

    // Re-map it if it's on its first chance.
    if (fte->referenced == 0) {
        frame_referenced(frame_number);
    }

    return FRAME_TABLE_SUCCESS;
}

// Releasing a frame = allowing it to be swapped again.
int frame_release(seL4_Word vaddr) {
    seL4_Word frame_number = frame_translate_vaddr(vaddr);
    frame_table_entry *fte = _frame_table + frame_number;

    // The frame needs to be pinned to be released
    if (fte->swappable) return FRAME_TABLE_INVAL;

    // Insert in the swap queue
    DL_APPEND(_swap_queue, fte);
    fte->referenced = 1;
    fte->swappable = 1;

    return FRAME_TABLE_SUCCESS;
}

// XXX clean up and add assertions
int frame_process_cap(seL4_Word sos_vaddr, sos_pcb *pcb, seL4_CPtr *proc_cap) {
    uint32_t frame_number = frame_translate_vaddr(sos_vaddr);
    frame_table_entry *fte = _frame_table + frame_number;
    // XXX Shared memory: use pcb to disambiguate
    assert(fte->mapping.proc_cap != 0);
    (*proc_cap) = fte->mapping.proc_cap;
    return FRAME_TABLE_SUCCESS;
}

// XXX proc_cspace argument is redundant
int frame_ungrant(seL4_Word vaddr, cspace_t *proc_cspace, seL4_CPtr cap) {
    int err;
    uint32_t frame_number;
    frame_table_entry *fte;

    dprintf(5, "frame_ungrant: %d @ %p\n", frame_translate_vaddr(vaddr), vaddr);

    // Check frame table initialised
    conditional_panic((_frame_table == NULL), "Frame table uninitialised");

    // Check cap exists
    if (cap == 0) {
        dprintf(6, "Attempt to ungrant null cap\n");
        return FRAME_TABLE_INVAL;
    }

    // Check frame is allocated and cap == fte->cap
    // XXX Deduplicate - standardise table indexing etc
    frame_number = frame_translate_vaddr(vaddr);
    fte = _vaddr_to_fte(vaddr);
    if (fte->page == 0 || fte->mapping.proc_cap != cap) {
        dprintf(6, "Frame not allocated / cap mismatch\n");
        return FRAME_TABLE_INVAL;
    }

    // Unmap if frame is still referenced
    if (fte->referenced) {
        err = seL4_ARM_Page_Unmap(cap);
        if (err) {
            dprintf(6, "frame_ungrant: page unmap failed\n");
            return FRAME_TABLE_OOM;
        }
    }

    // Delete cap
    err = cspace_delete_cap(cur_cspace, cap);
    conditional_panic(err, "Cspace failed to delete cap");

    if (fte->swappable) {
        // No longer swappable (if there are no process caps remaining)
        DL_DELETE(_swap_queue, fte);
        fte->swappable = 0;
    }

    // Clear the mapping
    memset((void *) &(fte->mapping), 0x00, sizeof(struct frame_mapping));

    dprintf(6, "frame_ungrant: finished\n");

    return FRAME_TABLE_SUCCESS;
}

inline void _fte_set_unmanaged(frame_table_entry *fte) {
    memset(fte, 0x00, sizeof(struct frame_table_entry));
}

int frame_free(seL4_Word vaddr) {
    int err;
    uint32_t frame_number;
    seL4_Word paddr;
    frame_table_entry *fte;

    // Die if uninitialised
    conditional_panic((_frame_table == NULL), "Frame table uninitialised");

    // Users should always provide valid vaddr within the window
    if (vaddr < WINDOW_START || vaddr > (WINDOW_END - (PAGE_SIZE - 1))) {
        return FRAME_TABLE_INVAL;
    }

    // We later convert the frame number into a page-aligned physical address
    paddr = PAGE_ALIGN(vaddr - WINDOW_START);
    frame_number = PAGE_SHIFT(paddr);
    fte = _frame_table + frame_number;

    // XXX Shared memory: check there are no more caps in the list.
    //     If caps remain, return separate code and exit here.
    // conditional_panic((fte->mapping.proc_cap != 0), "Attempt to free granted frame");
    if (fte->mapping.proc_cap != 0) {
        dprintf(6, "frame_free: freeing a granted frame\n");
        err = frame_ungrant(vaddr, cur_cspace, fte->mapping.proc_cap);
        if (err) return err;
    }

    // unmap from our window
    err = seL4_ARM_Page_Unmap(fte->page);
    if (err) {
        dprintf(6, "Failed to unmap page from SOS window\n");
        return err;
    }

    // Delete SmallPageObject cap from SOS cspace
    err = cspace_delete_cap(cur_cspace, fte->page);
    if (err) return err;

    // return the frame to ut
    ut_free(paddr, PAGE_BITS);

    // clear the table entry
    _fte_set_unmanaged(fte);

    // Reduce frame counter
    _cur_frames--;

    return FRAME_TABLE_SUCCESS;
}

seL4_Word frame_translate_vaddr(seL4_Word vaddr) {
    return (PAGE_SHIFT(vaddr - WINDOW_START));
}

seL4_Word frame_translate_number(seL4_Word frame_number) {
    return (WINDOW_START + PAGE_UNSHIFT(frame_number));
}

int frame_flush_icache(seL4_Word frame_number) {
    /* XXX Do bounds checking */
    frame_table_entry *fte = _frame_table + frame_number;
    dprintf(5, "frame_flush_icache: FN %d\tFTE %p\tcap %p\n",
            frame_number, fte, fte->page);
    seL4_CPtr cap = fte->page;
    seL4_ARM_Page_Unify_Instruction(cap, 0, PAGE_SIZE);
    return FRAME_TABLE_SUCCESS;
}

/* Map the (present) page back in for a process, and resets the
   referenced bit. This can be called even when a frame is already
   referenced, and should do nothing in that case. */
int frame_referenced(seL4_Word frame_number) {
    int err;
    frame_table_entry *fte = _frame_number_to_fte(frame_number);
    seL4_Word sos_vaddr = frame_translate_number(frame_number);
    // For each attached process,
    frame_mapping *map = &(fte->mapping);

    dprintf(6, "frame_referenced: FN %d, FTE %p, vaddr %p\n",
            frame_number, fte, sos_vaddr);

    if (fte->referenced == 0) {
        dprintf(5, "frame_referenced: mapping %d back in\n", frame_number);
        // map back in
        err = sos_page_map(map->pcb, map->proc_vaddr, map->proc_cap);
        // Toggle referenced bit
        fte->referenced = 1;
    } else {
        // Do nothing and return success if already mapped + referenced.
        err = FRAME_TABLE_SUCCESS;
    }
    return err;
}

/* The frame has had its first chance - unmap from all processes */
void _frame_chance(frame_table_entry *fte) {
    int err;

    dprintf(7, "frame_swapout_chance: first chance for fte %p\tfn %d\t%p\n",
            fte, _fte_to_framenumber(fte),
            frame_translate_number(_fte_to_framenumber(fte)));

    // For each attached process,
    frame_mapping *map = &(fte->mapping);

    // sos_page_unmap
    err = sos_page_unmap(map->pcb, map->proc_vaddr, map->proc_cap);

    // Mark referenced as 0
    fte->referenced = 0;
    return;
}

/* Retype and map a frame */
static seL4_Word _frame_retype_and_map(seL4_Word paddr, seL4_Word vaddr) {
    int err;
    seL4_Word cap;
    err = cspace_ut_retype_addr(paddr, seL4_ARM_SmallPageObject, PAGE_BITS,
                                cur_cspace, &cap);
    conditional_panic(err, "Failed to retype frame");

    err = map_page(cap, seL4_CapInitThreadPD, vaddr, seL4_AllRights,
                   seL4_ARM_Default_VMAttributes);
    conditional_panic(err, "Failed to map page into SOS");

    return cap;
}

/* Allocate enough pages for size bytes and map them contiguously from vaddr up */
/* (This is only to be used in table_init and leaks frames - beware) */
int _range_alloc_retype_and_map(seL4_Word vaddr, seL4_Word size) {
    int i;
    seL4_Word cur_vaddr = vaddr;
    seL4_Word num_frames = DIVROUND(size, PAGE_SIZE);
    int ret = FRAME_TABLE_SUCCESS;

    for(i = 0; i < num_frames; i++) {
        seL4_Word paddr = ut_alloc(PAGE_BITS);
        if (paddr) {
            _frame_retype_and_map(paddr, cur_vaddr);
            cur_vaddr += PAGE_SIZE;
        } else {
            dprintf(1, "paddr returned null\n");
            ret = FRAME_TABLE_OOM;
            break;
        }
    }
    dprintf(1, "alloced %d pages in range\n", i);

    return ret;
}

static void _test_indexing(void) {
    // Test indexing functions - they should produce the same FTE and
    // they should look sane w.r.t. frametable
    dprintf(1, "_frame_table: %p\tsizeof(fte): %d\n",
            _frame_table, sizeof(frame_table_entry));
    seL4_Word test_num = 0;
    while(test_num < 500) {
        frame_table_entry *fte = _frame_number_to_fte(test_num);
        seL4_Word test_vaddr = frame_translate_number(test_num);
        frame_table_entry *fte2 = _vaddr_to_fte(test_vaddr);

        dprintf(1, "FN: %d\tsos_vaddr: %p\tfte1: %p\tfte2: %p\n",
                test_num, test_vaddr, fte, fte2);

        test_num += 100;
    }
    return;
}

frame_table_entry *_frame_number_to_fte(seL4_Word frame_num) {
    return (_frame_table + frame_num);
}

seL4_Word _fte_to_framenumber(frame_table_entry *fte) {
    return (fte - _frame_table);
}

frame_table_entry *_vaddr_to_fte(seL4_Word sos_vaddr) {
    return (_frame_table + PAGE_SHIFT(sos_vaddr - WINDOW_START));
}
